package com.tao.springai.controller;

import com.tao.springai.service.ChatHistoryService;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

@RestController
@CrossOrigin(origins = "*")
public class MyController {

    private final ChatClient chatClient;
    private final ChatHistoryService chatHistoryService;

    public MyController(ChatClient.Builder builder, ChatHistoryService chatHistoryService) {
        // 给模型system一个角色，便于生成更准确的回答
        this.chatClient = builder.
                defaultSystem("你是一个智能机器人，精通各行业的知识。可以回答用户提出的任何问题")
                .build();
        this.chatHistoryService = chatHistoryService;
    }

    /**
     * 将 AI 模型的响应作为 String 返回。需等待模型全部响应完成后再返回，等待时间较久。
     * @param prompt 提示词
     * @return 模型的响应文本
     */
    @GetMapping("/chat/{prompt}")
    public String chat(@PathVariable("prompt") String prompt) {
        System.out.println("prompt:" + prompt);
        return chatClient.prompt()
                // 提问语
                .user(prompt)
                .call()
                .content();
    }

    /**
     * 流式响应API，前端可以逐步接收和显示AI生成的内容
     * @param prompt 提示词
     * @return 模型的流式响应文本
     */
    @GetMapping("/chat/stream/{prompt}")
    public Flux<String> chatStream(@PathVariable("prompt") String prompt) {
        System.out.println("Stream prompt:" + prompt);
        return chatClient.prompt()
                .user(prompt)
                .stream()
                .content();
    }

    /**
     * 在会话中进行聊天（携带会话ID）
     * @param sessionId 会话ID
     * @param requestBody 包含用户提示的请求体
     * @return 模型的响应文本
     */
    @PostMapping("/chat/session/{sessionId}")
    public String chatWithSession(
            @PathVariable String sessionId,
            @RequestBody Map<String, String> requestBody) {
        
        String prompt = requestBody.get("prompt");
        System.out.println("Session prompt: " + prompt);
        
        // 记录用户消息
        chatHistoryService.addMessage(sessionId, prompt, true);
        
        // 获取AI回复
        String response = chatClient.prompt()
                .user(prompt)
                .call()
                .content();
        
        // 记录AI回复
        chatHistoryService.addMessage(sessionId, response, false);
        
        return response;
    }

    /**
     * 在会话中进行流式聊天（携带会话ID）
     * @param sessionId 会话ID
     * @param requestBody 包含用户提示的请求体
     * @return 模型的流式响应文本
     */
    @PostMapping("/chat/stream/session/{sessionId}")
    public Flux<String> chatStreamWithSession(
            @PathVariable String sessionId,
            @RequestBody Map<String, String> requestBody) {
        
        String prompt = requestBody.get("prompt");
        
        System.out.println("Stream session prompt: " + prompt);
        
        // 记录用户消息
        chatHistoryService.addMessage(sessionId, prompt, true);
        
        // 获取并返回AI的流式回复
        Flux<String> responseFlux = chatClient.prompt()
                .user(prompt)
                .stream()
                .content();
        
        // 收集完整响应并保存到聊天历史
        StringBuilder fullResponse = new StringBuilder();
        
        return responseFlux
                .doOnNext(chunk -> {
                    fullResponse.append(chunk);
                })
                .doOnComplete(() -> {
                    // 流式传输完成后，保存完整的AI回复到聊天历史
                    chatHistoryService.addMessage(sessionId, fullResponse.toString(), false);
                    System.out.println("AI reply saved to history, length: " + fullResponse.length());
                });
    }
}


